"""
CTA策略回测模板
"""
import os
import json
import pandas as pd
import pickle
import math
from loguru import logger
from importlib import import_module
import subprocess
from yhlz_backtest.broker import FuturesBroker
from yhlz_backtest.feed import PandasFeed
from yhlz_backtest.plugins import ReturnsAnalyzer, TradeStatsAnalyzer, QuantstatsAnalyzer
from yhlz_backtest.ploter.streamlit_base import backtest as show_backtest


class CTATemplate:
    """
    CTA策略回测模板类
    
    用于加载配置、初始化策略、处理数据和分析回测结果
    """
    
    def __init__(self, config):
        """
        初始化CTA策略回测模板
        
        Parameters
        ----------
        config : dict
            配置字典，应包含以下字段：
            - strategy:
                - path: 策略模块路径
                - class: 策略类名
                - params: 策略参数字典（可选）
            - data:
                - dir: 数据目录路径
            - broker:
                - initial_capital: 初始资金
                - commission_rate: 手续费率
                - margin_rate: 保证金率
                - contract_multiplier: 合约乘数
                - slippage: 滑点设置
            - output:
                - dir: 输出目录路径
        """
        # 保存配置
        self.config = config
        
        # 确保输出目录存在
        os.makedirs(self.config['output']['dir'], exist_ok=True)
        
        # 初始化组件
        self.feed = PandasFeed()
        self.broker = FuturesBroker(
            initial_cash=self.config['broker']['initial_capital'],
            commission=self.config['broker']['commission_rate']
        )
        self.strategy = None
        self.analyzers = []
        
    def _init_feed(self):
        """初始化数据源"""
        # 加载并处理所有CSV文件
        for filename in os.listdir(self.config['data']['dir']):
            if filename.endswith('.csv'):
                # 从文件名获取品种代码
                instrument = os.path.splitext(filename)[0]
                filepath = os.path.join(self.config['data']['dir'], filename)
                
                # 读取CSV文件
                df = pd.read_csv(filepath)
                
                # 确保datetime列存在并转换为datetime类型
                if 'datetime' in df.columns:
                    df['datetime'] = pd.to_datetime(df['datetime'])
                    df.set_index('datetime', inplace=True)
                else:
                    logger.error(f"文件 {filename} 缺少datetime列")
                    continue
                
                # 调用策略的process_feed方法处理数据
                if self.strategy is not None:
                    processed_df = self.strategy.process_feed(df)
                    if processed_df is not None:
                        df = processed_df
                
                # 添加到数据源
                self.feed.add_dataframe(instrument, df)
        
        # 生成统一的时间轴
        self.feed.generate_timeline()
        
    def _init_broker(self):
        """初始化经纪人"""
        # 设置保证金率和合约乘数
        for instrument in self.feed.instruments:
            self.broker.set_margin_rate(instrument, self.config['broker']['margin_rate'])
            self.broker.set_multiplier(instrument, self.config['broker']['contract_multiplier'])
            self.broker.set_slippage(instrument, self.config['broker']['slippage'])
            
    def _init_strategy(self):
        """初始化策略"""
        try:
            # 动态导入策略模块
            strategy_module = import_module(self.config['strategy']['path'])
            strategy_class = getattr(strategy_module, self.config['strategy']['class'])
            
            # 实例化策略，如果有策略参数则传入
            if 'params' in self.config['strategy']:
                self.strategy = strategy_class(self.feed, self.broker, **self.config['strategy']['params'])
            else:
                self.strategy = strategy_class(self.feed, self.broker)
            
        except Exception as e:
            logger.error(f"初始化策略失败: {str(e)}")
            raise
            
    def _init_analyzers(self):
        """初始化分析器"""
        # 添加基础分析器
        self.analyzers.extend([
            ReturnsAnalyzer(self.broker.trades, self.broker.equity),
            TradeStatsAnalyzer(self.broker.trades, self.broker.equity),
            QuantstatsAnalyzer(self.broker.trades, self.broker.equity)
        ])
        
    def run(self):
        """
        运行回测
        
        Returns
        -------
        dict
            回测结果，包含各个分析器的分析结果
        """
        try:
            # 初始化策略
            logger.info("初始化策略...")
            self._init_strategy()
            
            # 初始化数据源
            logger.info("初始化数据源...")
            self._init_feed()
            
            # 初始化经纪人
            logger.info("初始化经纪人...")
            self._init_broker()
            
            # 运行回测
            logger.info("开始回测...")
            self.strategy.run()
            
            # 初始化分析器
            logger.info("初始化分析器...")
            self._init_analyzers()
            
            # 运行分析
            logger.info("分析回测结果...")
            results = {}
            for analyzer in self.analyzers:
                analyzer_name = analyzer.__class__.__name__
                results[analyzer_name] = analyzer.analyze()
                
                # 如果是QuantstatsAnalyzer，生成HTML报告
                if isinstance(analyzer, QuantstatsAnalyzer):
                    analyzer.save_html_report(
                        output_dir=self.config['output']['dir'],
                        filename='quantstats_report.html',
                        title=f"{self.config['strategy']['class']} 策略分析报告"
                    )
            
            logger.info("回测完成")
            return results
            
        except Exception as e:
            logger.error(f"回测过程中发生错误: {str(e)}")
            raise
            
    def _get_git_commit_id(self):
        """
        获取当前Git commit ID
        
        Returns
        -------
        str
            当前commit的短ID（前8位），如果获取失败则返回'unknown'
        """
        try:
            result = subprocess.run(
                ['git', 'rev-parse', '--short=8', 'HEAD'],
                capture_output=True,
                text=True,
                check=True
            )
            return result.stdout.strip()
        except Exception as e:
            logger.warning(f"获取Git commit ID失败: {str(e)}")
            return 'unknown'

    def _build_output_path(self):
        """
        构建输出路径
        
        根据commit ID和策略参数构建唯一的输出路径
        
        Returns
        -------
        str
            构建的输出路径
        """
        # 获取commit ID
        commit_id = self._get_git_commit_id()
        
        # 获取策略参数并排序，确保相同参数生成相同的路径
        if 'params' in self.config['strategy']:
            params = self.config['strategy']['params']
            # 将参数转换为排序后的字符串列表
            param_strs = []
            for key in sorted(params.keys()):
                value = params[key]
                # 如果值是数字，使用固定格式避免精度问题
                if isinstance(value, (int, float)):
                    param_str = f"{key}_{value:.6g}"
                else:
                    param_str = f"{key}_{value}"
                param_strs.append(param_str)
            
            # 将参数拼接成路径
            param_path = "_".join(param_strs)
        else:
            param_path = "default"
        
        # 构建最终路径：commit_id/params
        return os.path.join(self.config['output']['dir'], commit_id, param_path)
            
    def save_results(self, results):
        """
        保存回测结果
        
        Parameters
        ----------
        results : dict
            回测结果
        """
        try:
            # 构建输出路径
            output_path = self._build_output_path()
            
            # 确保输出目录存在
            os.makedirs(output_path, exist_ok=True)
            
            # 保存回测结果到JSON文件
            results_path = os.path.join(output_path, 'backtest_results.json')
            with open(results_path, 'w', encoding='utf-8') as f:
                # 添加策略参数到结果中
                results_with_params = {
                    'strategy_params': self.config['strategy'].get('params', {}),
                    'results': results
                }
                json.dump(results_with_params, f, indent=4, ensure_ascii=False)
            
            # 保存交易记录到CSV文件
            trades_path = os.path.join(output_path, 'trades.csv')
            self.broker.trades.to_csv(trades_path, index=True)
            
            # 保存权益曲线到CSV文件
            equity_path = os.path.join(output_path, 'equity.csv')
            self.broker.equity.to_csv(equity_path, index=True)
            
            # 如果有QuantstatsAnalyzer，保存HTML报告
            for analyzer in self.analyzers:
                if isinstance(analyzer, QuantstatsAnalyzer):
                    analyzer.save_html_report(
                        output_dir=output_path,
                        filename='quantstats_report.html',
                        title=f"{self.config['strategy']['class']} 策略分析报告"
                    )
            
            logger.info(f"回测结果已保存到目录: {output_path}")
            
            return output_path
            
        except Exception as e:
            logger.error(f"保存回测结果时发生错误: {str(e)}")
            raise

    def prepare_plot_data(self, output_path):
        """
        准备并保存可视化所需的数据
        
        Parameters
        ----------
        output_path : str
            输出目录路径
            
        Returns
        -------
        str
            输出目录路径
        """
        try:
            # 准备权益曲线数据
            equity_data = {
                'stg_name': self.config['strategy']['class'],
                'datetime': self.broker.equity.index.tolist(),
                'equity': self.broker.equity['equity'].tolist(),
                'equity_log': self.broker.equity['equity'].apply(lambda x: math.log10(x)).tolist(),
                'draw_down': self.broker.equity['drawdown'].tolist(),
                'equity_vs_market_value': self.broker.equity['market_value_ratio'].tolist(),
                'benchmark': self.broker.equity.get('benchmark', [0] * len(self.broker.equity)).tolist()
            }
            
            # 保存可视化数据
            with open(os.path.join(output_path, 'images_data.pkl'), 'wb') as f:
                pickle.dump(equity_data, f)
                
            # 保存因子数据（如果有）
            if hasattr(self.feed, 'get_all'):
                with open(os.path.join(output_path, 'factor_feed.pkl'), 'wb') as f:
                    pickle.dump(self.feed, f)
                    
            logger.info("可视化数据准备完成")
            return output_path
            
        except Exception as e:
            logger.error(f"准备可视化数据时发生错误: {str(e)}")
            raise


def run_backtest(config, show_result=True):
    """
    运行回测的便捷函数
    
    Parameters
    ----------
    config : dict
        配置字典，包含所有回测所需的配置信息
    show_result : bool, optional
        是否自动展示回测结果，默认为True
        
    Returns
    -------
    dict
        回测结果
    """
    # 创建回测实例
    backtest = CTATemplate(config)
    
    # 运行回测
    results = backtest.run()
    
    # 保存结果并准备可视化数据
    output_path = backtest.save_results(results)
    backtest.prepare_plot_data(output_path)
    
    # 如果需要，自动展示回测结果
    if show_result:
        show_backtest(output_path)
    
    return results


if __name__ == '__main__':
    # 示例用法
    config = {
        'strategy': {
            'path': 'yhlz_backtest.strategy.rb',
            'class': 'RBStrategy'
        },
        'data': {
            'dir': 'data/rb'
        },
        'broker': {
            'initial_capital': 1000000,
            'commission_rate': 0.0003,
            'margin_rate': 0.1,
            'contract_multiplier': 10,
            'slippage': 1.0
        },
        'output': {
            'dir': 'output/rb_strategy'
        }
    }
    results = run_backtest(config)
    print(json.dumps(results, indent=4, ensure_ascii=False)) 